use crate::{
    abort_on_drop::AbortOnDropHandle,
    net_ifce,
    proto::{
        BindingRequestPacket, BindingResponse, BindingResponsePacket, Header, HeaderPacket, MutableBindingIndicationPacket, MutableBindingRequestPacket, MutableBindingResponsePacket, MutableHeaderPacket, BINDING_INDICATION_LEN, BINDING_REQUEST_LEN, BINDING_RESPONSE_LEN, HEADER_LEN, PROTO_TYPE_BINDING_INDICATION, PROTO_TYPE_BINDING_REQUEST, PROTO_TYPE_BINDING_RESPONSE
    },
    raw_pcap, raw_udp,
    settings::{PeerEndpoint, Settings, DEFAULT_LISTEN_PORT},
};
use log;
use pnet::packet::{ipv4::Ipv4Packet, udp::UdpPacket};
use pnet_packet::{FromPacket, Packet, PacketSize};
use std::{
    collections::HashMap,
    net::{Ipv4Addr, SocketAddr, SocketAddrV4}, time::Duration,
};
use tokio::{net::UdpSocket, sync::mpsc};
use rand::{thread_rng, RngCore};

pub struct Server {
    handle: AbortOnDropHandle<()>,
    ctrl_tx: mpsc::Sender<ServerCtrlMsg>,
}
enum ServerCtrlMsg {
    Stop,
}

#[derive(Clone)]
enum ClientCtrlMsg {
    BindingResponse(Header, BindingResponse),
}

impl Server {
    pub fn new(cfg: &'static Settings) -> Self {
        let (tx, rx) = mpsc::channel(10);
        let handle = tokio::spawn(server_loop(cfg, tx.clone(), rx)).into();
        Self {
            handle,
            ctrl_tx: tx,
        }
    }
}

struct Client {
    state: ClientState,
    endpoint: PeerEndpoint,
    handle: AbortOnDropHandle<()>,
    ctrl_tx: mpsc::Sender<ClientCtrlMsg>,
}

enum ClientState {
    Init,
    Online,
}

async fn server_loop(
    cfg: &'static Settings,
    tx: mpsc::Sender<ServerCtrlMsg>,
    mut rx: mpsc::Receiver<ServerCtrlMsg>,
) {
    let mut client_map: HashMap<PeerEndpoint, Client> = HashMap::new();

    for ep in cfg.peer_endpoints.iter() {
        client_map.insert(ep.clone(), Client::new(cfg, ep.clone()));
    }

    let filter = format!("inbound and udp port {}", DEFAULT_LISTEN_PORT);
    let mut pcap = raw_pcap::RawPcapSocket::new(&filter);

    net_ifce::init_bridge(DEFAULT_LISTEN_PORT, cfg.virtual_ip);

    let udp_sock = raw_udp::RawUdpSocket::new();

    loop {
        tokio::select! {
            Some(msg) = rx.recv() => match msg {
                ServerCtrlMsg::Stop => break,
            },
            Some(pkt) = pcap.next_packet() => {
                log::info!("got len {} {:?}", pkt.len(), pkt);
                if pkt.len() >= 14 + 20 + 8 + HEADER_LEN {
                    let ip = Ipv4Packet::new(&pkt[14..]).unwrap();
                    let udp = UdpPacket::new(&pkt[(14 + 20)..]).unwrap();
                    let header = HeaderPacket::new(&pkt[(14 + 20 + 8)..]).unwrap();
                    if header.get_typ() == PROTO_TYPE_BINDING_REQUEST {
                        let req = BindingRequestPacket::new(&pkt[(14 + 20 + 8 + HEADER_LEN)..]).unwrap();
                        log::info!("recv binding request from {}:{}, {}", ip.get_source(), udp.get_source(), req.get_client_vip());
                        let mut buf = [0u8;1500];
                        let mut resp_header = MutableHeaderPacket::new(&mut buf).unwrap();
                        resp_header.set_typ(PROTO_TYPE_BINDING_RESPONSE);
                        resp_header.set_token(cfg.token);
                        resp_header.set_transaction(header.get_transaction());
                        drop(resp_header);
                        let mut resp_binding = MutableBindingResponsePacket::new(&mut buf).unwrap();
                        resp_binding.set_client_vip(req.get_client_vip());
                        resp_binding.set_server_vip(cfg.virtual_ip);
                        resp_binding.set_client_reflex_ip(ip.get_source());
                        resp_binding.set_client_reflex_port(udp.get_source());
                        drop(resp_binding);
                        let payload = &buf[..(HEADER_LEN + BINDING_RESPONSE_LEN)];

                        let dst_addr = SocketAddrV4::new(ip.get_source(), udp.get_source());
                        let src_addr = SocketAddrV4::new(ip.get_destination(), udp.get_destination());
                        udp_sock.send_pkt(src_addr, dst_addr, payload);

                    } else if header.get_typ() == PROTO_TYPE_BINDING_RESPONSE {
                        let resp = BindingResponsePacket::new(&pkt[(14 + 20 + 8 + HEADER_LEN)..]).unwrap();
                        let resp = resp.from_packet();
                        let msg = ClientCtrlMsg::BindingResponse(header.from_packet(), resp);
                        for client in client_map.values() {
                            let _ = client.ctrl_tx.send(msg.clone()).await;
                        }
                        log::info!("binding response {:?} {:?}", udp, header);
                    }
                }

            },
            else => break,
        }
    }
    log::warn!("server_loop exited");
}

impl Client {
    fn new(cfg: &'static Settings, endpoint: PeerEndpoint) -> Self {
        let (ctrl_tx, rx) = mpsc::channel(32);
        let handle = tokio::spawn(client_loop(cfg, endpoint.clone(), rx)).into();
        Self {
            state: ClientState::Init,
            endpoint,
            handle,
            ctrl_tx,
        }
    }
}

async fn client_loop(
    cfg: &'static Settings,
    endpoint: PeerEndpoint,
    mut rx: mpsc::Receiver<ClientCtrlMsg>,
) {
    let udp_sock = raw_udp::RawUdpSocket::new();
    let dst_addr = resolve_addr(&endpoint)
        .await
        .expect("resolve peer addr error");
    let src_addr = udp_sock
        .probe_local_addr(dst_addr)
        .expect("resolve local addr failed.");
    let mut ticker = tokio::time::interval(Duration::from_secs(5));
    let mut buf = [0u8; 1500];
    let mut txn_id = None;
    loop {
        tokio::select! {
            _ = ticker.tick() => {
                log::info!("send BINDING request from {} to {}", src_addr, dst_addr);
                let mut header = MutableHeaderPacket::new(&mut buf).unwrap();
                header.set_typ(PROTO_TYPE_BINDING_REQUEST);
                header.set_token(cfg.token);
                txn_id.replace(thread_rng().next_u32());
                header.set_transaction(txn_id.unwrap());
                drop(header);
        
                let mut req = MutableBindingRequestPacket::new(&mut buf[HEADER_LEN..]).unwrap();
                req.set_client_vip(cfg.virtual_ip);
                drop(req);
        
                let pkt = &buf[..(HEADER_LEN + BINDING_REQUEST_LEN)];
        
                udp_sock.send_pkt(src_addr, dst_addr, pkt);
            },
            Some(msg) = rx.recv() => match msg {
                ClientCtrlMsg::BindingResponse(resp_hdr, resp) => {
                    if Some(resp_hdr.transaction) == txn_id {

                        log::info!("binding response: {} -> {}", resp.client_vip, resp.server_vip);

                        log::info!("send BINDING indication from {} to {}", src_addr, dst_addr);
                        let mut header = MutableHeaderPacket::new(&mut buf).unwrap();
                        header.set_typ(PROTO_TYPE_BINDING_INDICATION);
                        header.set_token(cfg.token);
                        header.set_transaction(resp_hdr.transaction);
                        drop(header);
                
                        let mut ind = MutableBindingIndicationPacket::new(&mut buf[HEADER_LEN..]).unwrap();
                        ind.set_client_vip(resp.client_vip);
                        ind.set_server_vip(resp.server_vip);
                        ind.set_client_reflex_ip(resp.client_reflex_ip);
                        ind.set_client_reflex_port(resp.client_reflex_port);
                        drop(ind);
                
                        let pkt = &buf[..(HEADER_LEN + BINDING_INDICATION_LEN)];
                
                        udp_sock.send_pkt(src_addr, dst_addr, pkt);
                    }
                }
            },
            else => {
                break;
            }
        }
    }
}

async fn resolve_addr(endpoint: &PeerEndpoint) -> std::io::Result<SocketAddrV4> {
    match endpoint {
        PeerEndpoint::Host(host) => {
            for addr in tokio::net::lookup_host(host).await? {
                if let SocketAddr::V4(addr) = addr {
                    return Ok(addr);
                }
            }
            Err(std::io::Error::other(format!("resolve [{}] failed", host)))
        }
        PeerEndpoint::Addr(addr) => Ok(*addr),
    }
}
